package top.codetheory._05_二叉搜索树;

import java.util.Comparator;

/**
 * @author:
 * @since: 2023/11/28 14:05
 **/
public class BinarySearchTree<E> extends BinaryTree<E>  {

    private Comparator<E> comparator;

    public BinarySearchTree() {
        this(null);
    }

    public BinarySearchTree(Comparator<E> comparator) {
        this.comparator = comparator;
    }

    public void add(E element) {
        // element 不能为null
        elementNotNullCheck(element);

        // 添加第一个节点
        if (root == null) {
            root = createNode(element, null);  // 创建根节点
            size++;
            // 新添加节点之后的处理
            afterAdd(root);
            return;
        }

        // 添加的不是第一个节点
        // 找到父节点
        Node<E> parent = root;
        Node<E> node = root;
        int cmp = 0;

        do {
            cmp = compare(element, node.element);  //  比较两个元素的大小, 确定位置
            parent = node; // 记录parent
            if (cmp > 0) {
                node = node.right;  // 需要插入的元素比当前节点大, 那么往右找
            } else if (cmp < 0) {
                node = node.left;  // 需要插入的元素比当前节点小, 那么继续往左找
            } else {
                node.element = element;  // 相等的情况, 覆盖原有的值
                return;
            }
        } while (node != null);  // 一直往下找

        // 此时 parent 就是要插入元素的父节点, node 为待插入的位置, 但是node现在始终为null
        Node<E> newNode = createNode(element, parent);
        if (cmp > 0) { // 当前节点比父节点大, 那么就是插在右边
            parent.right = newNode;
        } else {
            parent.left = newNode;  // 比父节点小
        }
        size++;
        // 新添加节点之后的处理
        afterAdd(newNode);
    }
    public void remove(E element) {
        remove(getNodeByElement(element));
    }

    private void remove(Node<E> node) {
        if (node == null) {
            return;
        }
        size--;
        if (node.hasTwoChildren()) { // 度为2的节点
            // 找到后继节点
            Node<E> s = successor(node);
            // 用后继节点的值覆盖度为2的节点的值
            node.element = s.element;
            // 删除后继节点
            node = s;
        }
        // 删除node节点（node的度必然是1或者0）
        Node<E> replacement = node.left != null ? node.left : node.right;

        if (replacement != null) { // node是度为1的节点
            // 更改parent
            replacement.parent = node.parent;
            // 更改parent的left、right的指向
            if (node.parent == null) { // node是度为1的节点并且是根节点
                root = replacement;
            } else if (node == node.parent.left) {
                node.parent.left = replacement;
            } else { // node == node.parent.right
                node.parent.right = replacement;
            }
            // 删除节点之后的处理
//            afterRemove(node, replacement);
            afterRemove(replacement, null);
        } else if (node.parent == null) { // node是叶子节点并且是根节点
            root = null;
            // 删除节点之后的处理
            afterRemove(node, null);
        } else { // node是叶子节点，但不是根节点
            if (node == node.parent.left) {
                node.parent.left = null;
            } else { // node == node.parent.right
                node.parent.right = null;
            }
            // 删除节点之后的处理
            afterRemove(node, null);
        }

    }

    public boolean contains(E element) {
        return getNodeByElement(element) != null;
    }

    private Node<E> getNodeByElement(E element) {
        Node<E> node = root;
        while (node != null) {
            int cmp = compare(element, node.element);
            if (cmp == 0) return node;
            if (cmp > 0) {
                node = node.right;
            } else {
                node = node.left;
            }
        }
        return null;
    }

    /**
     * @return 返回值等于0，代表e1和e2相等；返回值大于0，代表e1大于e2；返回值小于于0，代表e1小于e2
     */
    private int compare(E e1, E e2) {
        if (comparator != null) {
            return comparator.compare(e1, e2);
        }
        return ((Comparable<E>)e1).compareTo(e2);  // e1 和 e2 是要可比较的, 要自己实现comparable接口, 否则会报错
    }

    private void elementNotNullCheck(E element) {
        if (element == null) {
            throw new IllegalArgumentException("元素不能为null");
        }
    }

    /**
     * 添加node之后的调整
     * @param node 新添加的节点
     */
    protected void afterAdd(Node<E> node) { }

    /**
     * 删除node之后的调整
     * @param node 被删除的节点
     */
    protected void afterRemove(Node<E> node, Node<E> replacement) { }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        toString(root, sb, "");
        return sb.toString();
    }
    private void toString(Node<E> node, StringBuilder sb, String prefix) {
        if (node == null) return;

        toString(node.left, sb, prefix + "L---");
        sb.append(prefix).append(node.element).append("\n");
        toString(node.right, sb, prefix + "R---");
    }


}